import requests
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from sklearn.preprocessing import MinMaxScaler, StandardScaler  # 新增StandardScaler用于对比不同归一化方式
from sklearn.model_selection import train_test_split  # 用于更灵活划分数据集
from sklearn.metrics import mean_absolute_error, mean_squared_error, mean_absolute_percentage_error  # 新增更多评估指标
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Bidirectional
import tushare as ts
import math
import random
import webbrowser

# 设置随机种子，保证结果的可复现性
random.seed(1)
np.random.seed(1)
tf.random.set_seed(1)

# 设置中文字体为黑体，解决负号显示问题（用于部分和matplotlib相关的操作，比如后续模型训练结果可视化等）
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 获取股票数据的函数，使用Tushare的数据接口（需替换token才能正常获取数据）
def get_stock_data(token, ts_code, start_date, end_date, fields):
    """
    从tushare获取指定股票代码、日期范围以及字段的股票数据

    参数:
    token (str): 在tushare平台注册获取的token
    ts_code (str): 股票代码，如 "01211.HK"
    start_date (str): 开始日期，格式为 "YYYYMMDD"
    end_date (str): 结束日期，格式为 "YYYYMMDD"
    fields (list): 要获取的股票数据字段列表，如 ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    返回:
    pd.DataFrame: 获取到的股票数据，如果获取失败返回None
    """
    try:
        pro = ts.pro_api(token)  # 请替换为自己的Tushare token
        data = pro.daily(ts_code=ts_code,
                         start_date=start_date,
                         end_date=end_date,
                         fields=fields)
        return data
    except Exception as e:
        print(f"获取股票数据时出现错误: {e}")
        return None


# 数据预处理函数，进行缺失值处理、数据归一化等操作，增加更多详细信息打印
def preprocess_data(data):
    df = data.copy()
    # 将交易日期列设置为索引，并转换为日期时间类型
    df.set_index('trade_date', inplace=True)
    df.index = pd.to_datetime(df.index)
    print("设置日期索引后的基本信息：")
    print(df.info())

    # 去除包含缺失值的行
    df.dropna(inplace=True)
    print("去除缺失值后的基本信息：")
    print(df.info())

    # 提取需要的列（开盘价、最高价、最低价、收盘价、成交量），可根据实际需求调整
    columns = ['open', 'high', 'low', 'close', 'vol']
    df = df[columns]
    print("提取关键列后的基本信息：")
    print(df.info())

    # 数据归一化，这里同时提供MinMaxScaler和StandardScaler两种归一化方式，可对比选择
    # 使用MinMaxScaler将数据归一化到0-1区间
    minmax_scaler = MinMaxScaler(feature_range=(0, 1))
    # 使用StandardScaler进行标准化（均值为0，标准差为1）
    standard_scaler = StandardScaler()

    df_minmax_scaled = minmax_scaler.fit_transform(df)
    df_standard_scaled = standard_scaler.fit_transform(df)
    print("MinMaxScaler归一化后的数据形状:", df_minmax_scaled.shape)
    print("StandardScaler归一化后的数据形状:", df_standard_scaled.shape)
    return df, df_minmax_scaled, df_standard_scaled, minmax_scaler, standard_scaler


# 创建训练数据集和测试数据集的函数，按照时间序列划分数据，使用sklearn的train_test_split更灵活划分
def create_dataset(dataset, time_steps=1):
    dataX, dataY = [], []
    for i in range(len(dataset) - time_steps - 1):
        a = dataset[i:(i + time_steps), :]
        dataX.append(a)
        dataY.append(dataset[i + time_steps, -1])
    return np.array(dataX), np.array(dataY)


# 构建并训练LSTM模型的函数，支持不同模型类型选择，添加超参数调优相关逻辑示例（简单演示网格搜索概念）
def build_and_train_model(model_type, X_train, y_train, X_test, y_test):
    # 简单示例超参数组合，实际可拓展更多组合进行细致调优
    units_choices = [32, 50, 64]
    epochs_choices = [20, 30, 40]
    batch_size_choices = [32, 64, 128]
    best_rmse = float('inf')
    best_model = None
    best_history = None
    for units in units_choices:
        for epochs in epochs_choices:
            for batch_size in batch_size_choices:
                model = Sequential()
                if model_type == 1:
                    # 单层 LSTM
                    model.add(LSTM(units=units, activation='relu',
                                   input_shape=(X_train.shape[1], X_train.shape[2])))
                    model.add(Dense(units=1))
                elif model_type == 2:
                    # 多层 LSTM
                    model.add(LSTM(units=units, activation='relu', return_sequences=True,
                                   input_shape=(X_train.shape[1], X_train.shape[2])))
                    model.add(LSTM(units=units, activation='relu'))
                    model.add(Dense(1))
                elif model_type == 3:
                    # 双向 LSTM
                    model.add(Bidirectional(LSTM(units, activation='relu'),
                                            input_shape=(X_train.shape[1], X_train.shape[2])))
                    model.add(Dense(1))
                model.compile(loss='mean_squared_error', optimizer='adam')
                history = model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=epochs,
                                    batch_size=batch_size, verbose=1)
                # 在测试集上评估模型，获取RMSE
                y_pred = model.predict(X_test)
                rmse = math.sqrt(mean_squared_error(y_test, y_pred))
                if rmse < best_rmse:
                    best_rmse = rmse
                    best_model = model
                    best_history = history
    return best_model, best_history


# 计算多个评估指标（均方根误差RMSE、平均绝对误差MAE、平均绝对百分比误差MAPE），用于更全面评估模型性能
def calculate_metrics(y_true, y_pred):
    # 确保 y_true 和 y_pred 的长度一致
    min_len = min(len(y_true), len(y_pred))  # 获取较小的长度
    y_true = y_true[:min_len]  # 修剪 y_true
    y_pred = y_pred[:min_len]  # 修剪 y_pred

    rmse = math.sqrt(mean_squared_error(y_true, y_pred))
    mae = mean_absolute_error(y_true, y_pred)
    mape = mean_absolute_percentage_error(y_true, y_pred)
    return rmse, mae, mape


if __name__ == "__main__":
    # Tushare平台获取的token，需替换为你自己的真实token
    token = '572b6125c826e36619a670a9943995412179dabcfb739c5c48913f46'
    ts_code = "600519.SH"  # 贵州茅台股票代码，可替换为其他感兴趣的股票代码
    start_date = "20100101"
    end_date = "20241225"
    fields = ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    # 获取股票数据
    stock_data = get_stock_data(token, ts_code, start_date, end_date, fields)

    # 检查获取的数据是否为空
    if stock_data is None or stock_data.empty:
        print("没有获取到有效的股票数据！")
        exit(1)

    # 数据预处理
    df, df_minmax_scaled, df_standard_scaled, minmax_scaler, standard_scaler = preprocess_data(stock_data)

    # 划分训练集和测试集，使用train_test_split，可指定test_size比例等更灵活
    test_size = 0.2  # 设定测试集占比为20%
    train_data_minmax, test_data_minmax = train_test_split(df_minmax_scaled, test_size=test_size, random_state=1)
    train_data_standard, test_data_standard = train_test_split(df_standard_scaled, test_size=test_size, random_state=1)
    time_steps = 10  # 时间步长
    X_train_minmax, y_train_minmax = create_dataset(train_data_minmax, time_steps)
    X_test_minmax, y_test_minmax = create_dataset(test_data_minmax, time_steps)
    X_train_standard, y_train_standard = create_dataset(train_data_standard, time_steps)
    X_test_standard, y_test_standard = create_dataset(test_data_standard, time_steps)
    X_train_minmax = np.reshape(X_train_minmax, (X_train_minmax.shape[0], X_train_minmax.shape[1], df.shape[1]))
    X_test_minmax = np.reshape(X_test_minmax, (X_test_minmax.shape[0], X_test_minmax.shape[1], df.shape[1]))
    X_train_standard = np.reshape(X_train_standard, (X_train_standard.shape[0], X_train_standard.shape[1], df.shape[1]))
    X_test_standard = np.reshape(X_test_standard, (X_test_standard.shape[0], X_test_standard.shape[1], df.shape[1]))

    # 选择模型类型（1: 单层 LSTM; 2: 多层 LSTM; 3: 双向 LSTM）
    model_type = 3
    # 使用MinMaxScaler归一化后的数据训练模型
    model_minmax, history_minmax = build_and_train_model(model_type, X_train_minmax, y_train_minmax, X_test_minmax,
                                                        y_test_minmax)
    # 使用StandardScaler归一化后的数据训练模型
    model_standard, history_standard = build_and_train_model(model_type, X_train_standard, y_train_standard,
                                                              X_test_standard, y_test_standard)

    # 绘制训练过程中的损失曲线（以MinMaxScaler对应的为例，同理可绘制StandardScaler对应的）
    plt.plot(history_minmax.history['loss'], label='训练损失（MinMaxScaler）')
    plt.plot(history_minmax.history['val_loss'], label='验证损失（MinMaxScaler）')
    plt.title('模型训练损失曲线（MinMaxScaler）')
    plt.xlabel('Epoch')
    plt.ylabel('损失值')
    plt.legend()
    plt.show()

    # 使用训练好的模型进行预测（以MinMaxScaler对应的为例，同理可对StandardScaler对应的操作）
    train_predict_minmax = model_minmax.predict(X_train_minmax)
    train_predict_minmax = np.reshape(train_predict_minmax, (train_predict_minmax.shape[0], 1))
    test_predict_minmax = model_minmax.predict(X_test_minmax)
    test_predict_minmax = np.reshape(test_predict_minmax, (test_predict_minmax.shape[0], 1))

    # 对训练集进行反归一化处理（MinMaxScaler）
    train_data_trimmed_minmax = train_data_minmax[time_steps:len(train_predict_minmax) + time_steps, 0:df.shape[1] - 1]
    train_predict_minmax = minmax_scaler.inverse_transform(np.concatenate((train_data_trimmed_minmax, train_predict_minmax),
                                                                         axis=1))[:, -1]
    # 对测试集进行反归一化处理（MinMaxScaler）
    test_data_trimmed_minmax = test_data_minmax[time_steps:len(test_predict_minmax) + time_steps, 0:df.shape[1] - 1]
    test_predict_minmax = minmax_scaler.inverse_transform(np.concatenate((test_data_trimmed_minmax, test_predict_minmax),
                                                                        axis=1))[:, -1]

    # 计算并打印多种评估指标（以MinMaxScaler对应的为例，同理可对StandardScaler对应的计算）
    rmse_minmax, mae_minmax, mape_minmax = calculate_metrics(df['close'].iloc[len(train_predict_minmax) + time_steps:len(
        train_predict_minmax) + len(test_predict_minmax) + time_steps - 1].values, test_predict_minmax)
    print("MinMaxScaler - 测试集均方根误差（RMSE）:", rmse_minmax)
    print("MinMaxScaler - 测试集平均绝对误差（MAE）:", mae_minmax)
    print("MinMaxScaler - 测试集平均绝对百分比误差（MAPE）:", mape_minmax)

    # 绘制股票价格走势与预测结果对比图（以MinMaxScaler对应的为例，同理可对StandardScaler对应的绘制）
    fig = go.Figure()

    # 添加真实股价走势（训练集）
    fig.add_trace(go.Scatter(x=df.index[time_steps:len(train_predict_minmax) + time_steps - 1],
                             y=df['close'].iloc[time_steps:len(train_predict_minmax) + time_steps - 1],
                             mode='lines',
                             name='真实股价（训练集）'))

    # 添加真实股价走势（测试集）
    fig.add_trace(go.Scatter(x=df.index[len(train_predict_minmax) + time_steps:len(train_predict_minmax) + len(
        test_predict_minmax) + time_steps - 1],
                             y=df['close'].iloc[len(train_predict_minmax) + time_steps:len(train_predict_minmax) + len(
                                 test_predict_minmax) + time_steps - 1],
                             mode='lines',
                             name='真实股价（测试集）'))

    # 添加预测股价走势（训练集）
    fig.add_trace(go.Scatter(x=df.index[time_steps:len(train_predict_minmax) + time_steps - 1],
                             y=train_predict_minmax,
                             mode='lines',
                             name='预测股价（训练集）',
                             line=dict(dash='dash')))

    # 添加预测股价走势（测试集）
    fig.add_trace(go.Scatter(x=df.index[len(train_predict_minmax) + time_steps:len(train_predict_minmax) + len(
        test_predict_minmax) + time_steps - 1],
                             y=test_predict_minmax,
                             mode='lines',
                             name='预测股价（测试集）',
                             line=dict(dash='dash')))

    # 设置图表标题和布局
    fig.update_layout(title='股票真实价格与预测价格对比（MinMaxScaler）',
                      xaxis_title='日期',
                      yaxis_title='股价',
                      legend=dict(x=0.8, y=0.9),  # 调整图例位置
                      hovermode='x unified')  # 悬停交互模式

    # 显示图表
    fig.show()